/*
 * EncDecrypterTests.cpp
 *
 *  Created on: Aug 6, 2010
 *      Author: Lukasz Spintzyk
 */

#include "CudaAESInvOperations.h"
#include "CudaAESOperations.h"
#include "CudaAESFast.h"
#include "cudaUtils.h"
#include "CUDAAESEncrypter.h"
#include "OpenSSLAESEncrypter.h"
#include "BigInteger.h"

#include "aesUtils.h"
#include "AesKeyTestVectors.h"
#include <iostream>
#include <iomanip>
#include <string>
#include <ctime>
#include <cstdio>
#include <cstdlib>
#include <string>

#include <cppunit/TestRunner.h>
#include <cppunit/TestResult.h>
#include <cppunit/TestResultCollector.h>
#include <cppunit/extensions/HelperMacros.h>
#include <cppunit/BriefTestProgressListener.h>
#include <cppunit/extensions/TestFactoryRegistry.h>

#include <openssl/aes.h>

byte myKey[32] =
{ 123, 23, 45, 66, 21, 10, 05, 87, 33, 62, 92, 11, 0, 1, 78, 51, 9, 200, 234,
		2, 19, 203, 34, 12, 9, 200, 234, 2, 12, 243, 37, 77,

};
class TestEncDecrypter: public CPPUNIT_NS::TestCase
{
CPPUNIT_TEST_SUITE(TestEncDecrypter);
		CPPUNIT_TEST(testIVEncryption);
		CPPUNIT_TEST(testIVEncryption2048);
		CPPUNIT_TEST(testIVEncryption3072);
		CPPUNIT_TEST(testIVEncryption4x1024);
		CPPUNIT_TEST(testIVEncryptionN);
	CPPUNIT_TEST_SUITE_END();

public:
	void setUp(void)
	{
		size = 1024 * 15;
		host_data = new byte[size];
		host_data2 = new byte[size];
		hostBuffer = new byte[size];

		for (size_t i = 0; i < size; ++i)
		{
			host_data[i] = host_data2[i] = i;
		}
		size = 1024;
		//vect = "11223344556677889900AABBCCDDEEFC";
		//vect =   "11111111111111111111111111111111";
		vect =   "FF000000000000000000000000000000";
		p_iv = new BigInteger(vect.c_str(), 0);
		p_encrypter = new CUDAAESEncrypter(*p_iv, myKey);
		p_openSSLEnc = new OpenSSLAESEncrypter(*p_iv, myKey);

		//aesEncryptInit();
		//aesDecryptInit();
		//aesFastCipherInit();
	}
	void tearDown(void)
	{
		delete[] host_data;
		delete[] host_data2;
		delete[] hostBuffer;
		delete p_encrypter;
	}

	void testIVEncryption()
	{

		std::string vect = "01223344556677889900AABBCCDDEEFF";
		BigInteger iv(vect.c_str(), 0);
		CUDAAESEncrypter encrypter(iv, myKey);
		memset(host_data, 0, size);

		byte* outPtr;
		byte* outPtr2;
		try
		{
			outPtr = (byte*) encrypter.encrypt(host_data, size, 0);
		}
		catch(CUDAException& e)
		{
			std::cout << "CUDA Excrption " << e.what() << std::endl;
		}
		//std::cout << std::endl;
		//printBlock((byte*)outPtr, 70, 4);


		/*		byte ivb[16]={0};
		 iv.getBinValue(ivb);
		 //printBlock(ivb, 4,4);
		 AES_encrypt(ivb, ivb, &encrypter.getKey());
		 std::cout << std::endl;
		 printBlock(ivb, 4,4);
		 */
		memset(host_data2, 0, size);
		OpenSSLAESEncrypter openSSLEnc(iv, myKey);
		//for (int i =0; i < 100; i++)
		outPtr2 = (byte*) openSSLEnc.encrypt(host_data2, 1024, 0);

		//std::cout << std::endl;
		//printBlock(outPtr2, 70, 4);


		for (uint32_t i = 0; i < size; ++i)
		{
			if (host_data[i] != host_data2[i])
			{
				//std::cout << "Error "<<i<< " "<< (unsigned int)host_data[i] << " "<< (unsigned int)host_data2[i] <<std::endl;
			}

			CPPUNIT_ASSERT(outPtr[i] == outPtr2[i]);
		}
	}

	void testIVEncryption2048()
	{
		size = 2048;
		memset(host_data, 0, size);
		memset(host_data2, 0, size);

		byte* outPtr;
		byte* outPtr2;
		outPtr = (byte*) p_encrypter->encrypt(host_data, size, 0);
		//printBlock((byte*) outPtr, 70, 4);

		outPtr2 = (byte*) p_openSSLEnc->encrypt(host_data2, size, 0);
		//printBlock(outPtr2, 70, 4);

		for (uint32_t i = 0; i < size; ++i)
		{
			if (outPtr[i] != outPtr2[i])
			{
				std::cout << std::endl;
				std::cout << "Error " << i << " "
						<< (unsigned int) outPtr[i] << " "
						<< (unsigned int) outPtr2[i] << std::endl;
			}

			CPPUNIT_ASSERT(outPtr[i] == outPtr2[i]);
		}
	}

	void testIVEncryption3072()
	{
		size = 1024*3;
		memset(host_data, 0, size);
		memset(host_data2, 0, size);

		byte* outPtr;
		byte* outPtr2;
		outPtr = (byte*) p_encrypter->encrypt(host_data, size, 0);
		//printBlock((byte*) outPtr + 1024*2, 70, 4);


		outPtr2 = (byte*) p_openSSLEnc->encrypt(host_data2, size, 0);
		//printBlock(outPtr2 +1024*2, 70, 4);


		for (uint32_t i = 0; i < size; ++i)
		{
			if (outPtr[i] != outPtr2[i])
			{
				std::cout << std::endl;
				std::cout << "Error " << i << " "
						<< (unsigned int) outPtr[i] << " "
						<< (unsigned int) outPtr2[i] << std::endl;
			}

			CPPUNIT_ASSERT(outPtr[i] == outPtr2[i]);
		}
	}

	void testIVEncryption4x1024()
	{
		size = 1024*4;
		memset(host_data, 0, size);
		memset(host_data2, 0, size);

		byte* outPtr;
		byte* outPtr2;
		outPtr = (byte*) p_encrypter->encrypt(host_data, size, 0);

		outPtr2 = (byte*) p_openSSLEnc->encrypt(host_data2, size, 0);



		for (uint32_t i = 0; i < size; ++i)
		{
			if (outPtr[i] != outPtr2[i])
			{
				std::cout << std::endl;
				std::cout << "\nError " << i << " "
						<< (unsigned int) outPtr[i] << " "
						<< (unsigned int) outPtr2[i] << std::endl;

				printBlock(outPtr +i, 70, 4);
				printBlock(outPtr2 + i, 70, 4);
				//return;
			}

			CPPUNIT_ASSERT(outPtr[i] == outPtr2[i]);
		}
	}

	void testIVEncryptionN()
	{
		testIVDecryptionNx1024(13);
		for (int i = 2; i < 15; ++i)
		{
			//std::cout << i << std::endl;
		//	testIVEncryptionNx1024(5);
			//testIVEncryptionNx1024(i);
			testIVDecryptionNx1024(i);
		}
	}

	void testIVDecryptionNx1024(int n)
		{
			size = 1024*n;
			//memset(host_data, 0, size);
			//memset(host_data2, 0, size);
			host_data[123]= 0xFC;

			byte* outPtr;
			byte* outPtr2;
			outPtr = (byte*) p_encrypter->encrypt(host_data, size, 0);
			memcpy(hostBuffer, outPtr, size);
			outPtr = (byte*) p_encrypter->encrypt(hostBuffer, size, 0);
			//printBlock(host_data , 70, 4);
			//printBlock(outPtr , 70, 4);
			//printBlock(hostBuffer , 70, 4);
/*			CPPUNIT_ASSERT(memcmp(outPtr, host_data, size) == 0);

			for (uint32_t i = 0; i < size; ++i)
				{
					if (outPtr[i] != host_data[i])
					{
						std::cout << std::endl;
						std::cout << "Error " << i << " "
								<< (unsigned int) outPtr2[i] << " "
								<< (unsigned int) host_data[i] << std::endl;

						printBlock(outPtr +i, 70, 4);
						printBlock(host_data + i, 70, 4);
						return;
					}
				}
*/
			outPtr2 = (byte*) p_openSSLEnc->encrypt(host_data, size, 0);
			memcpy(host_data2, outPtr2, size);
			outPtr2 = (byte*) p_openSSLEnc->encrypt(host_data2, size, 0);



			//CPPUNIT_ASSERT(memcmp(outPtr, host_data, size) == 0);
			CPPUNIT_ASSERT(memcmp(outPtr2, host_data, size) == 0);
			for (uint32_t i = 0; i < size; ++i)
			{
				if (outPtr[i] != host_data[i])
				{
					std::cout << std::endl;
					std::cout << "Error " << i << " "
							<< (unsigned int) outPtr2[i] << " "
							<< (unsigned int) host_data[i] << std::endl;

					printBlock(outPtr +i, 70, 4);
					printBlock(host_data + i, 70, 4);

				}

				CPPUNIT_ASSERT(outPtr[i] == host_data[i]);
			}
		}



	void testIVEncryptionNx1024(int n)
	{
		size = 1024*n;
		//memset(host_data, 0, size);
		//memset(host_data2, 0, size);

		byte* outPtr;
		byte* outPtr2;
		outPtr = (byte*) p_encrypter->encrypt(host_data, size, 0);
		//memcpy(hostBuffer, outPtr, size);
		//outPtr = (byte*) p_encrypter->decrypt(hostBuffer, size, 0);



		outPtr2 = (byte*) p_openSSLEnc->encrypt(host_data2, size, 0);
 		 byte ivb[16]={0};
 		memcpy(ivb, p_iv->value->d, 16);
		 AES_encrypt(ivb, ivb, &p_encrypter->getKey());
		 //std::cout << std::endl;

		//	printBlock(outPtr , 70, 4);
		//	printBlock(outPtr2 , 70, 4);
		//	printBlock(ivb, 4,4);

		for (uint32_t i = 0; i < size; ++i)
		{
			if (outPtr[i] != outPtr2[i])
			{
				std::cout << std::endl;
				std::cout << "Error " << i << " "
						<< (unsigned int) outPtr[i] << " "
						<< (unsigned int) outPtr2[i] << std::endl;

				printBlock(outPtr +i, 70, 4);
				printBlock(outPtr2 + i, 70, 4);

			}

			CPPUNIT_ASSERT(outPtr[i] == outPtr2[i]);
		}
	}

protected:
	void printBlock(byte* block, int w, int h)
	{
		std::cout << std::endl;
		for (int j = 0; j < h; ++j)
		{
			for (int i = 0; i < w; ++i)
			{
				if (block[i * h + j] == 0xe4)
				{
					//std::cout << "[";
				}
				std::cout << std::hex << std::setw(3) << (unsigned int) block[i
						* h + j];
				if (block[i * h + j] == 0xe4)
				{
					//std::cout << "]";
				}
			}
			std::cout << std::endl;
		}
	}

private:

	byte* host_data;
	byte* host_data2;
	byte* hostBuffer;
	uint32_t size;


	std::string vect;
	BigInteger* p_iv;
	CUDAAESEncrypter* p_encrypter;
	OpenSSLAESEncrypter* p_openSSLEnc;

};
CPPUNIT_TEST_SUITE_REGISTRATION(TestEncDecrypter);

int main(int ac, char **av)
{
	//--- Create the event manager and test controller
	CPPUNIT_NS::TestResult controller;
	//--- Add a listener that colllects test result
	CPPUNIT_NS::TestResultCollector result;
	controller.addListener(&result);
	//--- Add a listener that print dots as test run.
	CPPUNIT_NS::BriefTestProgressListener progress;
	controller.addListener(&progress);
	//--- Add the top suite to the test runner
	CPPUNIT_NS::TestRunner runner;
	runner.addTest(CPPUNIT_NS::TestFactoryRegistry::getRegistry().makeTest());
	aesFastCipherInit();
	runner.run(controller);
	return result.wasSuccessful() ? 0 : 1;
}

